package org.springframework.social.security.provider;

import java.util.HashSet;
import java.util.Optional;
import java.util.Set;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.social.connect.Connection;
import org.springframework.social.connect.ConnectionFactory;
import org.springframework.social.connect.support.OAuth2ConnectionFactory;
import org.springframework.social.oauth2.AccessGrant;
import org.springframework.social.oauth2.OAuth2Parameters;
import org.springframework.social.security.SocialAuthenticationRedirectException;
import org.springframework.social.security.SocialAuthenticationToken;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestClientException;

public class OAuth2AuthenticationService<S> extends AbstractSocialAuthenticationService<S> {
  protected final Log logger = LogFactory.getLog(getClass());
  
  private OAuth2ConnectionFactory<S> connectionFactory;
  
  private Set<String> returnToUrlParameters;
  
  private String defaultScope = "";
  
  public OAuth2AuthenticationService(OAuth2ConnectionFactory<S> connectionFactory) {
    setConnectionFactory(connectionFactory);
  }
  
  public OAuth2ConnectionFactory<S> getConnectionFactory() {
    return this.connectionFactory;
  }
  
  public void setConnectionFactory(OAuth2ConnectionFactory<S> connectionFactory) {
    this.connectionFactory = connectionFactory;
  }
  
  public void setReturnToUrlParameters(Set<String> returnToUrlParameters) {
    Assert.notNull(returnToUrlParameters, "returnToUrlParameters cannot be null");
    this.returnToUrlParameters = returnToUrlParameters;
  }
  
  public Set<String> getReturnToUrlParameters() {
    if (this.returnToUrlParameters == null)
      this.returnToUrlParameters = new HashSet<>(); 
    this.returnToUrlParameters.add("channel");
    return this.returnToUrlParameters;
  }
  
  public void setDefaultScope(String defaultScope) {
    this.defaultScope = defaultScope;
  }
  
  public void afterPropertiesSet() throws Exception {
    super.afterPropertiesSet();
    Assert.notNull(getConnectionFactory(), "connectionFactory");
  }
  
  public SocialAuthenticationToken getAuthToken(HttpServletRequest request, HttpServletResponse response) throws SocialAuthenticationRedirectException {
    String code = request.getParameter("code");
    if (!StringUtils.hasText(code)) {
      OAuth2Parameters params = new OAuth2Parameters();
      params.setRedirectUri(buildReturnToUrl(request));
      setScope(request, params);
      params.add("state", generateState(this.connectionFactory, request));
      throw new SocialAuthenticationRedirectException(getConnectionFactory().getOAuthOperations().buildAuthenticateUrl(params));
    } 
    if (StringUtils.hasText(code))
      try {
        String returnToUrl = buildReturnToUrl(request);
        AccessGrant accessGrant = getConnectionFactory().getOAuthOperations().exchangeForAccess(code, returnToUrl, null);
        Connection<S> connection = getConnectionFactory().createConnection(accessGrant);
        return new SocialAuthenticationToken(connection, null);
      } catch (RestClientException e) {
        this.logger.debug("failed to exchange for access", (Throwable)e);
        return null;
      }  
    return null;
  }
  
  private String generateState(OAuth2ConnectionFactory<?> connectionFactory, HttpServletRequest request) {
    String state = Optional.<Object>ofNullable(request.getAttribute("state")).map(o -> (String)o).orElse(request.getParameter("state"));
    return (state != null) ? state : connectionFactory.generateState();
  }
  
  protected String buildReturnToUrl(HttpServletRequest request) {
    StringBuffer sb = request.getRequestURL();
    if (enableHttps(request) && sb.indexOf("http://") == 0)
      sb.replace(0, "http://".length(), "https://"); 
    if (!sb.toString().endsWith("callback"))
      if (sb.toString().endsWith("/")) {
        sb.append("callback");
      } else {
        sb.append("/").append("callback");
      }  
    sb.append("?");
    for (String name : getReturnToUrlParameters()) {
      String value = request.getParameter(name);
      if (value == null)
        continue; 
      sb.append(name).append("=").append(value).append("&");
    } 
    sb.setLength(sb.length() - 1);
    return sb.toString();
  }
  
  private boolean enableHttps(HttpServletRequest request) {
    return ((Boolean)request.getAttribute("enableHttps")).booleanValue();
  }
  
  private void setScope(HttpServletRequest request, OAuth2Parameters params) {
    String requestedScope = request.getParameter("scope");
    if (StringUtils.hasLength(requestedScope)) {
      params.setScope(requestedScope);
    } else if (StringUtils.hasLength(this.defaultScope)) {
      params.setScope(this.defaultScope);
    } 
  }
}
